{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#需要使用L4GPU，然后官方代码的专家数，从64改为16\n",
        "#把依赖文件kernel.py上传\n",
        "#注意本地没有GPU没法运行这个代码，要用colab或者autodl，torch版本太低不可以，因为没有rms归一化，还有本地没安装triton不可以"
      ],
      "metadata": {
        "id": "GuaaCdL-T76S"
      },
      "id": "GuaaCdL-T76S",
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip list|grep triton\n",
        "# 要求结果是triton 3.1.0"
      ],
      "metadata": {
        "id": "aNIvj2QZUS-C"
      },
      "id": "aNIvj2QZUS-C",
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Zq9JCGw0UokB",
        "outputId": "01b9f4d2-7513-4d54-ac2f-6b5a1d0e1d2f"
      },
      "id": "Zq9JCGw0UokB",
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "kernel.py  __pycache__\tsample_data\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "790dd8f7-da52-45fe-91ca-b2f4f613da6f",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "790dd8f7-da52-45fe-91ca-b2f4f613da6f",
        "outputId": "97194a8e-2f7e-4710-d785-5ee26714e5fa"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.5.1+cu124\n"
          ]
        }
      ],
      "source": [
        "import math\n",
        "from dataclasses import dataclass\n",
        "from typing import Tuple, Optional, Literal\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributed as dist\n",
        "\n",
        "\n",
        "print(torch.__version__)\n",
        "\n",
        "from kernel import act_quant, weight_dequant, fp8_gemm\n",
        "\n",
        "# 定义全局变量\n",
        "world_size = 1  # 世界大小（用于分布式计算）\n",
        "rank = 0  # 当前进程的排名\n",
        "block_size = 128  # 块大小（可能用于矩阵运算优化）\n",
        "gemm_impl: Literal[\"bf16\", \"fp8\"] = \"bf16\"  # 矩阵乘法实现方式（bf16 或 fp8）\n",
        "attn_impl: Literal[\"naive\", \"absorb\"] = \"absorb\"  # 注意力机制实现方式\n",
        "\n",
        "@dataclass\n",
        "class ModelArgs:\n",
        "    \"\"\"\n",
        "    用于定义模型参数和超参数的数据类。\n",
        "\n",
        "    属性：\n",
        "        max_batch_size (int): 最大批量大小。\n",
        "        max_seq_len (int): 最大序列长度。\n",
        "        dtype (Literal[\"bf16\", \"fp8\"]): 计算数据类型（bf16 或 fp8）。\n",
        "        vocab_size (int): 词汇表大小。\n",
        "        dim (int): 模型隐藏层维度。\n",
        "        inter_dim (int): MLP 层的中间层维度。\n",
        "        moe_inter_dim (int): MoE 层的中间层维度。\n",
        "        n_layers (int): Transformer 层的数量。\n",
        "        n_dense_layers (int): 模型中的全连接层数量。\n",
        "        n_heads (int): 注意力头的数量。\n",
        "        n_routed_experts (int): MoE 层中可路由的专家数量。\n",
        "        n_shared_experts (int): MoE 层中共享的专家数量。\n",
        "        n_activated_experts (int): MoE 层中每次激活的专家数量。\n",
        "        n_expert_groups (int): MoE 层中的专家组数量。\n",
        "        n_limited_groups (int): MoE 路由中的受限组数量。\n",
        "        score_func (Literal[\"softmax\", \"sigmoid\"]): MoE 路由的评分函数。\n",
        "        route_scale (float): MoE 路由评分的缩放因子。\n",
        "        q_lora_rank (int): 查询（Query）投影的 LoRA 秩。\n",
        "        kv_lora_rank (int): 键值（Key-Value）投影的 LoRA 秩。\n",
        "        qk_nope_head_dim (int): 无位置编码的 Query-Key 投影维度。\n",
        "        qk_rope_head_dim (int): 使用旋转位置编码（RoPE）的 Query-Key 投影维度。\n",
        "        v_head_dim (int): 值（Value）投影维度。\n",
        "        original_seq_len (int): 原始序列长度。\n",
        "        rope_theta (float): 旋转位置编码的基数。\n",
        "        rope_factor (float): 扩展序列长度的缩放因子。\n",
        "        beta_fast (int): 快速 Beta 修正因子。\n",
        "        beta_slow (int): 慢速 Beta 修正因子。\n",
        "        mscale (float): 扩展注意力的缩放因子。\n",
        "    \"\"\"\n",
        "    max_batch_size: int = 8  # 训练时的最大批次大小\n",
        "    max_seq_len: int = 4096 * 4  # 允许的最大序列长度（可能用于扩展序列训练）\n",
        "    dtype: Literal[\"bf16\", \"fp8\"] = \"bf16\"  # 计算数据类型，默认使用 bfloat16\n",
        "    vocab_size: int = 102400  # 词汇表大小\n",
        "    dim: int = 2048  # 模型隐藏层维度\n",
        "    inter_dim: int = 10944  # MLP 层的中间层维度\n",
        "    moe_inter_dim: int = 1408  # MoE 层的中间层维度\n",
        "    n_layers: int = 27  # Transformer 层数\n",
        "    n_dense_layers: int = 1  # 全连接层数（可能用于 MoE 或 MLP 结构）\n",
        "    n_heads: int = 16  # 注意力头数\n",
        "\n",
        "    # MoE（Mixture of Experts）相关参数\n",
        "    n_routed_experts: int = 16  # 可被路由的专家数量,论文里DeepSeekV3用了256个专家\n",
        "    n_shared_experts: int = 2  # 共享专家数量，始终都激活的专家数量，来保证模型的基线性能\n",
        "    n_activated_experts: int = 6  # 每次激活的专家数量\n",
        "    n_expert_groups: int = 1  # 专家组数量（可能用于分组专家路由）\n",
        "    n_limited_groups: int = 1  # 受限专家组数量\n",
        "    score_func: Literal[\"softmax\", \"sigmoid\"] = \"softmax\"  # MoE 路由评分函数，默认使用 softmax\n",
        "    route_scale: float = 1.0  # 路由评分的缩放因子\n",
        "\n",
        "    # MLA（Multi-Level Attention）相关参数\n",
        "    q_lora_rank: int = 0  # Query 投影的 LoRA 秩（低秩适配）\n",
        "    kv_lora_rank: int = 512  # Key-Value 投影的 LoRA 秩\n",
        "    qk_nope_head_dim: int = 128  # 无 RoPE 的 Query-Key 维度\n",
        "    qk_rope_head_dim: int = 64  # 使用 RoPE 的 Query-Key 维度\n",
        "    v_head_dim: int = 128  # Value 维度\n",
        "\n",
        "    # YARN（Yet Another RoPE Network）相关参数\n",
        "    original_seq_len: int = 4096  # 原始序列长度\n",
        "    rope_theta: float = 10000.0  # RoPE 旋转位置编码的基数\n",
        "    rope_factor: float = 40  # 扩展序列长度的缩放因子\n",
        "    beta_fast: int = 32  # 快速 Beta 修正因子\n",
        "    beta_slow: int = 1  # 慢速 Beta 修正因子\n",
        "    mscale: float = 1.0  # 扩展注意力的缩放因子"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "cec5b9fc4b12cf08",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-02-08T12:32:34.592983Z",
          "start_time": "2025-02-08T12:32:34.585392Z"
        },
        "id": "cec5b9fc4b12cf08"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributed as dist\n",
        "from typing import Optional\n",
        "\n",
        "# 假设 world_size 和 rank 由分布式训练环境提供\n",
        "world_size = dist.get_world_size() if dist.is_initialized() else 1\n",
        "rank = dist.get_rank() if dist.is_initialized() else 0\n",
        "\n",
        "class ParallelEmbedding(nn.Module):\n",
        "    \"\"\"\n",
        "    并行嵌入层（ParallelEmbedding），支持分布式训练环境下的词向量分片。\n",
        "\n",
        "    参数:\n",
        "        vocab_size (int): 词表大小，即整个模型的词汇总数。\n",
        "        dim (int): 词向量的维度。\n",
        "\n",
        "    说明:\n",
        "        - 词表在多个进程（GPU）之间进行分片，每个进程仅存储词表的一部分。\n",
        "        - vocab_size 必须能够被 world_size 整除，以确保各 GPU 拥有相同大小的词向量片段。\n",
        "    \"\"\"\n",
        "    def __init__(self, vocab_size: int, dim: int):\n",
        "        super().__init__()\n",
        "        self.vocab_size = vocab_size  # 词表大小\n",
        "        self.dim = dim  # 词向量维度\n",
        "        assert vocab_size % world_size == 0, f\"词表大小必须能被 world_size 整除 (world_size={world_size})\"\n",
        "\n",
        "        # 计算当前进程（GPU）负责的词表片段大小\n",
        "        self.part_vocab_size = vocab_size // world_size\n",
        "        # 计算当前进程的词表起始和结束索引\n",
        "        self.vocab_start_idx = rank * self.part_vocab_size\n",
        "        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size\n",
        "        # 初始化当前进程的词向量参数\n",
        "        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        并行嵌入层的前向传播。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量，包含 token 的索引。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 词向量表示。\n",
        "\n",
        "        处理流程:\n",
        "            1. 如果使用多 GPU 训练（world_size > 1），检查 token 是否属于当前 GPU 负责的词表范围:\n",
        "                - 若 token 超出当前 GPU 负责的范围，设为 0（避免索引超界）。\n",
        "            2. 计算词嵌入。\n",
        "            3. 若是多 GPU 训练，则进行 all_reduce 操作，将所有 GPU 的嵌入求和（同步）。\n",
        "        \"\"\"\n",
        "        if world_size > 1:\n",
        "            # 生成掩码，标记不属于当前进程词表范围的 token\n",
        "            mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)\n",
        "            # 将输入索引映射到当前进程的词表范围内\n",
        "            x = x - self.vocab_start_idx\n",
        "            # 将超出范围的索引设为 0，避免索引超界\n",
        "            x[mask] = 0\n",
        "\n",
        "        # 获取嵌入\n",
        "        y = F.embedding(x, self.weight)\n",
        "\n",
        "        if world_size > 1:\n",
        "            # 对超出范围的 token 设置为 0\n",
        "            y[mask] = 0\n",
        "            # all_reduce 操作，确保所有进程得到相同的词嵌入\n",
        "            dist.all_reduce(y)\n",
        "\n",
        "        return y\n",
        "\n",
        "\n",
        "def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    线性变换函数，实现 y = xA^T + b，支持量化权重的计算。\n",
        "\n",
        "    参数:\n",
        "        x (torch.Tensor): 输入张量。\n",
        "        weight (torch.Tensor): 权重张量，可能是量化后的，需要进行解量化处理。\n",
        "        bias (Optional[torch.Tensor]): 偏置项（可选），默认为 None。\n",
        "\n",
        "    返回:\n",
        "        torch.Tensor: 线性变换后的结果。\n",
        "\n",
        "    说明:\n",
        "        - 若 weight 不是量化的，则直接调用 F.linear 计算。\n",
        "        - 若 weight 是量化的（element_size() == 1），则需要先进行解量化，再进行计算。\n",
        "        - 当 gemm_impl == \"bf16\" 时，使用 bf16 计算。\n",
        "        - 其他情况，对 x 进行量化，然后使用 fp8_gemm 计算。\n",
        "    \"\"\"\n",
        "    if weight.element_size() > 1:\n",
        "        # 直接使用标准的 F.linear 计算\n",
        "        return F.linear(x, weight, bias)\n",
        "    elif gemm_impl == \"bf16\":\n",
        "        # 量化权重，需要解量化\n",
        "        weight = weight_dequant(weight, weight.scale)\n",
        "        return F.linear(x, weight, bias)\n",
        "    else:\n",
        "        # 其他情况：对 x 进行量化，并使用 fp8_gemm 计算\n",
        "        x, scale = act_quant(x, block_size)\n",
        "        y = fp8_gemm(x, scale, weight, weight.scale)\n",
        "        if bias is not None:\n",
        "            y += bias\n",
        "        return y\n",
        "\n",
        "\n",
        "class Linear(nn.Module):\n",
        "    \"\"\"\n",
        "    自定义线性层，支持量化权重，并提供可选的偏置项。\n",
        "\n",
        "    参数:\n",
        "        in_features (int): 输入特征维度。\n",
        "        out_features (int): 输出特征维度。\n",
        "        bias (bool): 是否包含偏置项，默认为 False。\n",
        "        dtype (可选): 计算数据类型，默认为 torch.bfloat16。\n",
        "\n",
        "    说明:\n",
        "        - 如果 weight 是量化的，则需要额外存储 scale 参数。\n",
        "        - bias 可选，若不使用，则注册为 None。\n",
        "    \"\"\"\n",
        "    dtype = torch.bfloat16  # 默认数据类型\n",
        "\n",
        "    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):\n",
        "        super().__init__()\n",
        "        self.in_features = in_features  # 输入特征维度\n",
        "        self.out_features = out_features  # 输出特征维度\n",
        "        self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype))\n",
        "\n",
        "        # 若权重是量化的（element_size() == 1），则需要 scale 参数\n",
        "        if self.weight.element_size() == 1:\n",
        "            scale_out_features = (out_features + block_size - 1) // block_size\n",
        "            scale_in_features = (in_features + block_size - 1) // block_size\n",
        "            # 存储量化 scale 参数\n",
        "            self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float32))\n",
        "        else:\n",
        "            # 非量化情况，无需 scale 参数\n",
        "            self.register_parameter(\"scale\", None)\n",
        "\n",
        "        # 处理偏置\n",
        "        if bias:\n",
        "            self.bias = nn.Parameter(torch.empty(out_features))\n",
        "        else:\n",
        "            self.register_parameter(\"bias\", None)\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        线性层的前向传播。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 经过线性变换的张量。\n",
        "\n",
        "        说明:\n",
        "            - 调用 linear 函数，自动处理量化权重和偏置项的计算。\n",
        "        \"\"\"\n",
        "        return linear(x, self.weight, self.bias)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "47a7753016e61fe3",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-02-08T12:45:32.907956Z",
          "start_time": "2025-02-08T12:45:32.903284Z"
        },
        "id": "47a7753016e61fe3"
      },
      "outputs": [],
      "source": [
        "class ColumnParallelLinear(Linear):\n",
        "    \"\"\"\n",
        "    列并行线性层（Column Parallel Linear），将输出特征分割到多个分布式进程中。\n",
        "\n",
        "    参数：\n",
        "        in_features (int): 输入特征的数量。\n",
        "        out_features (int): 总输出特征数量。\n",
        "        bias (bool): 是否包含偏置项，默认为 False。\n",
        "        dtype (optional): 数据类型，默认为 `torch.bfloat16`。\n",
        "    \"\"\"\n",
        "    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):\n",
        "        # 确保总输出特征数量可以被世界大小整除，以实现均匀分割\n",
        "        assert out_features % world_size == 0, f\"输出特征数必须能被 world_size 整除 (world_size={world_size})\"\n",
        "\n",
        "        # 计算当前进程负责的部分输出特征数\n",
        "        self.part_out_features = out_features // world_size\n",
        "\n",
        "        # 调用父类 Linear 的初始化，创建一个 in_features 到 part_out_features 的线性层\n",
        "        super().__init__(in_features, self.part_out_features, bias, dtype)\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        列并行线性层的前向传播。\n",
        "\n",
        "        参数：\n",
        "            x (torch.Tensor): 输入张量。\n",
        "\n",
        "        返回：\n",
        "            torch.Tensor: 经过线性变换后的张量，进行列并行计算。\n",
        "        \"\"\"\n",
        "        # 进行线性变换\n",
        "        y = linear(x, self.weight, self.bias)\n",
        "        return y\n",
        "\n",
        "\n",
        "class RowParallelLinear(Linear):\n",
        "    \"\"\"\n",
        "    行并行线性层（Row Parallel Linear），将输入特征分割到多个分布式进程中。\n",
        "\n",
        "    参数：\n",
        "        in_features (int): 总输入特征数量。\n",
        "        out_features (int): 输出特征的数量。\n",
        "        bias (bool): 是否包含偏置项，默认为 False。\n",
        "        dtype (optional): 数据类型，默认为 `torch.bfloat16`。\n",
        "    \"\"\"\n",
        "    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):\n",
        "        # 确保输入特征数量可以被世界大小整除，以实现均匀分割\n",
        "        assert in_features % world_size == 0, f\"输入特征数必须能被 world_size 整除 (world_size={world_size})\"\n",
        "\n",
        "        # 计算当前进程负责的部分输入特征数\n",
        "        self.part_in_features = in_features // world_size\n",
        "\n",
        "        # 调用父类 Linear 的初始化，创建一个 part_in_features 到 out_features 的线性层\n",
        "        super().__init__(self.part_in_features, out_features, bias, dtype)\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        行并行线性层的前向传播。\n",
        "\n",
        "        参数：\n",
        "            x (torch.Tensor): 输入张量。\n",
        "\n",
        "        返回：\n",
        "            torch.Tensor: 经过线性变换后的张量，进行行并行计算。\n",
        "        \"\"\"\n",
        "        # 进行线性变换\n",
        "        y = linear(x, self.weight)\n",
        "\n",
        "        # 如果是分布式环境（world_size > 1），则对 y 进行 all_reduce 操作，使所有进程的计算结果进行累加\n",
        "        if world_size > 1:\n",
        "            dist.all_reduce(y)\n",
        "\n",
        "        # 如果存在偏置项，则加上偏置\n",
        "        if self.bias is not None:\n",
        "            y += self.bias\n",
        "\n",
        "        return y\n",
        "\n",
        "\n",
        "class RMSNorm(nn.Module):\n",
        "    \"\"\"\n",
        "    均方根归一化（RMSNorm），用于对输入张量进行归一化。\n",
        "\n",
        "    该方法不同于标准 LayerNorm，不依赖均值，而是基于均方根（RMS）进行归一化。\n",
        "\n",
        "    参数：\n",
        "        dim (int): 输入张量的维度。\n",
        "        eps (float): 用于数值稳定性的 epsilon 值，默认为 1e-6。\n",
        "    \"\"\"\n",
        "    def __init__(self, dim: int, eps: float = 1e-6):\n",
        "        super().__init__()\n",
        "        self.dim = dim  # 记录输入维度\n",
        "        self.eps = eps  # 记录 epsilon 值\n",
        "\n",
        "        # 归一化的可训练缩放参数，初始化为全 1\n",
        "        self.weight = nn.Parameter(torch.ones(dim))\n",
        "\n",
        "    def forward(self, x: torch.Tensor):\n",
        "        \"\"\"\n",
        "        均方根归一化的前向传播。\n",
        "\n",
        "        参数：\n",
        "            x (torch.Tensor): 输入张量。\n",
        "\n",
        "        返回：\n",
        "            torch.Tensor: 归一化后的张量，保持输入形状不变。\n",
        "        \"\"\"\n",
        "        # 使用 torch 的 rms_norm 进行均方根归一化\n",
        "        return F.rms_norm(x, (self.dim,), self.weight, self.eps)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "3d07b29bba86a874",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-02-08T12:50:23.848766Z",
          "start_time": "2025-02-08T12:50:23.826915Z"
        },
        "id": "3d07b29bba86a874"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import math\n",
        "from typing import Optional\n",
        "\n",
        "# 预计算旋转位置编码的频率值\n",
        "# 该函数用于计算基于旋转位置编码的复指数值\n",
        "# 主要目的是为了加速计算，避免在每次前向传播时重新计算这些值\n",
        "def precompute_freqs_cis(args: ModelArgs) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    预计算旋转位置编码的频率值。\n",
        "\n",
        "    参数：\n",
        "        args (ModelArgs): 包含位置编码参数的模型参数。\n",
        "\n",
        "    返回：\n",
        "        torch.Tensor: 预计算的复指数值，用于旋转位置编码。\n",
        "    \"\"\"\n",
        "    dim = args.qk_rope_head_dim  # 旋转位置编码的维度,是64\n",
        "    seqlen = args.max_seq_len  # 最大序列长度\n",
        "    beta_fast = args.beta_fast  # 快速调整参数\n",
        "    beta_slow = args.beta_slow  # 缓慢调整参数\n",
        "    base = args.rope_theta  # 旋转位置编码的基数\n",
        "    factor = args.rope_factor  # 旋转位置编码的缩放因子\n",
        "\n",
        "    # 计算旋转位置编码修正维度\n",
        "    def find_correction_dim(num_rotations, dim, base, max_seq_len):\n",
        "        \"\"\"\n",
        "        计算旋转位置编码的修正维度。\n",
        "        \"\"\"\n",
        "        return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))\n",
        "\n",
        "    # 计算旋转位置编码修正范围\n",
        "    def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):\n",
        "        \"\"\"\n",
        "        计算旋转位置编码修正范围。\n",
        "        \"\"\"\n",
        "        low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))\n",
        "        high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))\n",
        "        return max(low, 0), min(high, dim-1)\n",
        "\n",
        "    # 计算线性斜坡因子，用于平滑过渡\n",
        "    def linear_ramp_factor(min, max, dim):\n",
        "        \"\"\"\n",
        "        计算线性斜坡因子。\n",
        "        \"\"\"\n",
        "        if min == max:\n",
        "            max += 0.001\n",
        "        linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)\n",
        "        ramp_func = torch.clamp(linear_func, 0, 1)\n",
        "        return ramp_func\n",
        "\n",
        "    # 计算基础频率\n",
        "    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))\n",
        "\n",
        "    # 如果序列长度超过原始最大长度，则进行修正\n",
        "    if seqlen > args.original_seq_len:\n",
        "        low, high = find_correction_range(beta_fast, beta_slow, dim, base, args.original_seq_len)\n",
        "        smooth = 1 - linear_ramp_factor(low, high, dim // 2)\n",
        "        freqs = freqs / factor * (1 - smooth) + freqs * smooth\n",
        "\n",
        "    t = torch.arange(seqlen)\n",
        "    freqs = torch.outer(t, freqs)\n",
        "    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)\n",
        "    return freqs_cis\n",
        "\n",
        "# 应用旋转位置编码到输入张量\n",
        "# 该函数使用预计算的复指数值对输入进行旋转编码\n",
        "\n",
        "def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    应用旋转位置编码到输入张量。\n",
        "\n",
        "    参数：\n",
        "        x (torch.Tensor): 输入张量。\n",
        "        freqs_cis (torch.Tensor): 预计算的复指数值。\n",
        "\n",
        "    返回：\n",
        "        torch.Tensor: 旋转编码后的张量。\n",
        "    \"\"\"\n",
        "    dtype = x.dtype\n",
        "    x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))  # 变换为复数表示\n",
        "    freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))  # 调整形状\n",
        "    y = torch.view_as_real(x * freqs_cis).flatten(3)  # 进行旋转编码并转换回实数\n",
        "    return y.to(dtype)\n",
        "\n",
        "# 多头注意力层（MLA）\n",
        "# 该类实现了标准的多头注意力机制，并结合了旋转位置编码\n",
        "class MLA(nn.Module):\n",
        "    \"\"\"\n",
        "    多头注意力层（MLA）。\n",
        "\n",
        "    属性:\n",
        "        dim (int): 输入特征的维度。\n",
        "        n_heads (int): 注意力头的数量。\n",
        "        n_local_heads (int): 分布式系统中用于局部注意力的头数量。\n",
        "        q_lora_rank (int): 查询低秩投影的秩。\n",
        "        kv_lora_rank (int): 键值低秩投影的秩。\n",
        "        qk_nope_head_dim (int): 非位置查询/键投影的维度。\n",
        "        qk_rope_head_dim (int): 旋转位置查询/键投影的维度。\n",
        "        qk_head_dim (int): 查询/键投影的总维度。\n",
        "        v_head_dim (int): 值投影的维度。\n",
        "        softmax_scale (float): 注意力计算中Softmax的缩放因子。\n",
        "    \"\"\"\n",
        "    def __init__(self, args: ModelArgs):\n",
        "        super().__init__()\n",
        "        # 初始化各个参数\n",
        "        self.dim = args.dim  # 输入的特征维度\n",
        "        self.n_heads = args.n_heads  # 注意力头的数量\n",
        "        self.n_local_heads = args.n_heads // world_size  # 分布式环境中的局部注意力头数\n",
        "        self.q_lora_rank = args.q_lora_rank  # 查询低秩投影的秩\n",
        "        self.kv_lora_rank = args.kv_lora_rank  # 键值低秩投影的秩\n",
        "        self.qk_nope_head_dim = args.qk_nope_head_dim  # 非位置查询/键的维度\n",
        "        self.qk_rope_head_dim = args.qk_rope_head_dim  # 旋转位置查询/键的维度\n",
        "        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim  # 查询/键投影的总维度\n",
        "        self.v_head_dim = args.v_head_dim  # 值投影的维度\n",
        "\n",
        "        # 如果q_lora_rank为0，直接使用列并行线性层，否则使用低秩投影和标准化\n",
        "        if self.q_lora_rank == 0:\n",
        "            self.wq = ColumnParallelLinear(self.dim, self.n_heads * self.qk_head_dim)\n",
        "        else:\n",
        "            self.wq_a = Linear(self.dim, self.q_lora_rank)  # 低秩投影的第一部分，不为0，为512，先从2048变为512\n",
        "            self.q_norm = RMSNorm(self.q_lora_rank)  # 低秩投影的标准化\n",
        "            self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)  # 低秩投影的第二部分\n",
        "\n",
        "        # 键值投影和标准化\n",
        "        self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)  # 键值低秩投影\n",
        "        self.kv_norm = RMSNorm(self.kv_lora_rank)  # 键值的标准化\n",
        "        self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))  # 键值投影\n",
        "        self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)  # 输出投影\n",
        "        self.softmax_scale = self.qk_head_dim ** -0.5  # Softmax缩放因子\n",
        "\n",
        "        # 如果最大序列长度大于原始序列长度，调整softmax_scale\n",
        "        if args.max_seq_len > args.original_seq_len:\n",
        "            mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0\n",
        "            self.softmax_scale = self.softmax_scale * mscale * mscale\n",
        "\n",
        "        # 根据注意力实现类型选择不同的缓存方式\n",
        "        if attn_impl == \"naive\":\n",
        "            # 在\"naive\"实现下缓存k和v\n",
        "            self.register_buffer(\"k_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim), persistent=False)\n",
        "            self.register_buffer(\"v_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim), persistent=False)\n",
        "        else:\n",
        "            # 在其他实现下缓存kv和pe\n",
        "            self.register_buffer(\"kv_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank), persistent=False)\n",
        "            self.register_buffer(\"pe_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim), persistent=False)\n",
        "\n",
        "    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):\n",
        "        \"\"\"\n",
        "        多头注意力层的前向传播。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量，形状为(batch_size, seq_len, dim)。\n",
        "            start_pos (int): 缓存的起始位置。\n",
        "            freqs_cis (torch.Tensor): 预计算的旋转嵌入的复数指数值。\n",
        "            mask (Optional[torch.Tensor]): 掩码张量，用于排除某些位置的注意力计算。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 输出张量，形状与输入相同。\n",
        "        \"\"\"\n",
        "        bsz, seqlen, _ = x.size()  # 获取输入张量的batch size和序列长度\n",
        "        end_pos = start_pos + seqlen  # 计算结束位置\n",
        "\n",
        "        # 计算查询（q）\n",
        "        if self.q_lora_rank == 0:\n",
        "            q = self.wq(x)  # 使用列并行线性层计算查询\n",
        "        else:\n",
        "            q = self.wq_b(self.q_norm(self.wq_a(x)))  # 使用低秩投影和标准化计算查询\n",
        "\n",
        "        # 重塑查询张量的形状\n",
        "        q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)\n",
        "        #使用f格式打印q的形状\n",
        "        print(f\"q shape: {q.shape}\")\n",
        "        # 分割查询张量为非位置部分和位置编码部分\n",
        "        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)\n",
        "        #使用f格式打印q_nope, q_pe的形状，把密集向量分成两部分，q_pe计算旋转位置编码\n",
        "        print(f\"q_nope shape: {q_nope.shape}, q_pe shape: {q_pe.shape}\")\n",
        "        q_pe = apply_rotary_emb(q_pe, freqs_cis)  # 应用旋转位置编码\n",
        "\n",
        "        # 计算键值（kv）\n",
        "        kv = self.wkv_a(x)\n",
        "        #打印kv shape\n",
        "        print(\"kv shape:\", kv.shape)\n",
        "        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)\n",
        "        #一行代码f格式打印kv，k_pe shape\n",
        "        print(f\"kv shape: {kv.shape}, k_pe shape: {k_pe.shape}\")\n",
        "        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)  # 应用旋转位置编码\n",
        "\n",
        "        # 判断是否使用\"naive\"注意力实现\n",
        "        if attn_impl == \"naive\":\n",
        "            q = torch.cat([q_nope, q_pe], dim=-1)  # 拼接查询的非位置部分和位置编码部分\n",
        "            kv = self.wkv_b(self.kv_norm(kv))  # 对键值进行标准化\n",
        "            kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)  # 重塑键值张量形状\n",
        "            # 分割键值张量为非位置部分和值（v）\n",
        "            k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)\n",
        "            k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)  # 拼接键的非位置部分和位置编码部分\n",
        "            #一行代码f格式打印k shape\n",
        "            print(f\"k shape: {k.shape}\")\n",
        "            self.k_cache[:bsz, start_pos:end_pos] = k  # 更新k缓存，是推理阶段\n",
        "            self.v_cache[:bsz, start_pos:end_pos] = v  # 更新v缓存\n",
        "            # 计算注意力得分\n",
        "            scores = torch.einsum(\"bshd,bthd->bsht\", q, self.k_cache[:bsz, :end_pos]) * self.softmax_scale\n",
        "        else:\n",
        "            # 在其他实现下，处理wkv_b的权重，并计算q_nope的注意力得分\n",
        "            wkv_b = self.wkv_b.weight if self.wkv_b.scale is None else weight_dequant(self.wkv_b.weight, self.wkv_b.scale, block_size)\n",
        "            wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)\n",
        "            q_nope = torch.einsum(\"bshd,hdc->bshc\", q_nope, wkv_b[:, :self.qk_nope_head_dim])\n",
        "            self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv)  # 更新kv缓存\n",
        "            self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)  # 更新位置编码缓存\n",
        "            # 计算注意力得分\n",
        "            scores = (torch.einsum(\"bshc,btc->bsht\", q_nope, self.kv_cache[:bsz, :end_pos]) +\n",
        "                      torch.einsum(\"bshr,btr->bsht\", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale\n",
        "\n",
        "        # 如果有mask，添加到得分上\n",
        "        if mask is not None:\n",
        "            scores += mask.unsqueeze(1)\n",
        "\n",
        "        # 计算softmax得分\n",
        "        scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x)\n",
        "\n",
        "        # 根据注意力实现类型，选择不同的计算方式\n",
        "        if attn_impl == \"naive\":\n",
        "            x = torch.einsum(\"bsht,bthd->bshd\", scores, self.v_cache[:bsz, :end_pos])  # 计算输出\n",
        "        else:\n",
        "            x = torch.einsum(\"bsht,btc->bshc\", scores, self.kv_cache[:bsz, :end_pos])  # 计算输出\n",
        "            x = torch.einsum(\"bshc,hdc->bshd\", x, wkv_b[:, -self.v_head_dim:])  # 计算最终输出\n",
        "\n",
        "        # 通过行并行线性层投影到原始维度\n",
        "        x = self.wo(x.flatten(2))\n",
        "        return x\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "e70257e6400c2689",
      "metadata": {
        "id": "e70257e6400c2689"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from typing import Tuple\n",
        "\n",
        "class MLP(nn.Module):\n",
        "    \"\"\"\n",
        "    多层感知机（MLP），用于前馈计算。\n",
        "\n",
        "    该模块包含三个线性变换层，分别是 w1、w2 和 w3，用于特征变换和计算。\n",
        "\n",
        "    属性:\n",
        "        w1 (nn.Module): 线性层，用于从输入层到隐藏层的转换。\n",
        "        w2 (nn.Module): 线性层，用于从隐藏层到输出层的转换。\n",
        "        w3 (nn.Module): 额外的线性层，用于特征变换。\n",
        "    \"\"\"\n",
        "    def __init__(self, dim: int, inter_dim: int):\n",
        "        \"\"\"\n",
        "        初始化 MLP 层。\n",
        "\n",
        "        参数:\n",
        "            dim (int): 输入和输出的维度（维度保持一致）。\n",
        "            inter_dim (int): 隐藏层的维度。\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.w1 = ColumnParallelLinear(dim, inter_dim)  # 第一层线性变换\n",
        "        self.w2 = RowParallelLinear(inter_dim, dim)     # 第二层线性变换（回到原始维度）\n",
        "        self.w3 = ColumnParallelLinear(dim, inter_dim)  # 额外的线性变换层\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        MLP 的前向计算。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量，形状为 (batch_size, dim)。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 经过 MLP 计算后的输出张量，形状为 (batch_size, dim)。\n",
        "        \"\"\"\n",
        "        return self.w2(F.silu(self.w1(x)) * self.w3(x))  # 使用 SiLU 激活函数进行非线性变换，并结合 w3(x) 进行特征变换\n",
        "\n",
        "\n",
        "class Gate(nn.Module):\n",
        "    \"\"\"\n",
        "    用于 Mixture-of-Experts（MoE）模型的门控机制（Gating Mechanism）。\n",
        "\n",
        "    该模块用于在多个专家（Expert）之间进行路由选择，决定每个输入数据应该被送到哪些专家进行计算。\n",
        "\n",
        "    属性:\n",
        "        dim (int): 输入特征的维度。\n",
        "        topk (int): 每个输入激活的专家数（选择 top-k 个专家）。\n",
        "        n_groups (int): 专家被划分的组数（用于路由分组）。\n",
        "        topk_groups (int): 每个输入路由到的专家组数。\n",
        "        score_func (str): 计算分数的函数（可选 \"softmax\" 或 \"sigmoid\"）。\n",
        "        route_scale (float): 路由权重的缩放因子。\n",
        "        weight (torch.nn.Parameter): 可训练参数，表示门控网络的权重矩阵。\n",
        "        bias (Optional[torch.nn.Parameter]): 可选的偏置项，仅当输入维度为 7168 时存在。\n",
        "    \"\"\"\n",
        "    def __init__(self, args: ModelArgs):\n",
        "        \"\"\"\n",
        "        初始化 Gate 模块。\n",
        "\n",
        "        参数:\n",
        "            args (ModelArgs): 传入的模型参数对象，包含 MoE 相关的超参数。\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.dim = args.dim  # 输入特征的维度\n",
        "        self.topk = args.n_activated_experts  # 选择的前 top-k 个专家\n",
        "        self.n_groups = args.n_expert_groups  # 专家分组的数量\n",
        "        self.topk_groups = args.n_limited_groups  # 选择的前 top-k 组\n",
        "        self.score_func = args.score_func  # 计算分数的方式\n",
        "        self.route_scale = args.route_scale  # 路由权重的缩放比例\n",
        "\n",
        "        # 可训练权重参数（用于计算门控分数）\n",
        "        self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))\n",
        "\n",
        "        # 只有当 dim = 7168 时，才会添加可训练的偏置项\n",
        "        self.bias = nn.Parameter(torch.empty(args.n_routed_experts)) if self.dim == 7168 else None\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
        "        \"\"\"\n",
        "        计算门控权重，并确定选择哪些专家进行计算。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量，形状为 (batch_size, dim)。\n",
        "\n",
        "        返回:\n",
        "            Tuple[torch.Tensor, torch.Tensor]:\n",
        "                - 选择的专家权重 (batch_size, topk)\n",
        "                - 选择的专家索引 (batch_size, topk)\n",
        "        \"\"\"\n",
        "        # 一行代码f格式打印x的形状\n",
        "        print(f\"Gate x shape: {x.shape}\")\n",
        "        scores = linear(x, self.weight)  # 计算输入与门控权重的线性变换\n",
        "\n",
        "        # 根据 score_func 选择 softmax 或 sigmoid 进行归一化\n",
        "        if self.score_func == \"softmax\":\n",
        "            scores = scores.softmax(dim=-1, dtype=torch.float32)\n",
        "        else:\n",
        "            scores = scores.sigmoid()\n",
        "\n",
        "        original_scores = scores  # 保存原始分数\n",
        "\n",
        "        # 若存在偏置项，则加上偏置\n",
        "        if self.bias is not None:\n",
        "            scores = scores + self.bias\n",
        "\n",
        "        # 若使用多个专家组，则进行分组处理\n",
        "        if self.n_groups > 1:\n",
        "            scores = scores.view(x.size(0), self.n_groups, -1)  # 重新 reshape 为 (batch_size, n_groups, 每组的专家数)\n",
        "\n",
        "            # 计算每组的得分，若无偏置，则取最大值；若有偏置，则取 top-2 得分之和\n",
        "            if self.bias is None:\n",
        "                group_scores = scores.amax(dim=-1)\n",
        "            else:\n",
        "                group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)\n",
        "\n",
        "            # 选择得分最高的 topk_groups 组，并生成掩码\n",
        "            indices = group_scores.topk(self.topk_groups, dim=-1)[1]\n",
        "            mask = torch.zeros_like(scores[..., 0]).scatter_(1, indices, True)\n",
        "            scores = (scores * mask.unsqueeze(-1)).flatten(1)  # 仅保留选中的专家分数\n",
        "\n",
        "        # 选择得分最高的 top-k 个专家\n",
        "        indices = torch.topk(scores, self.topk, dim=-1)[1]\n",
        "        #打印indices\n",
        "        print(\"Gate indices shape:\", indices.shape)\n",
        "        # 计算最终的专家权重\n",
        "        weights = original_scores.gather(1, indices)\n",
        "        #打印weights\n",
        "        # print(\"Gate weights :\", weights)\n",
        "        # 若使用 sigmoid，则需要归一化\n",
        "        if self.score_func == \"sigmoid\":\n",
        "            weights /= weights.sum(dim=-1, keepdim=True)\n",
        "\n",
        "        weights *= self.route_scale  # 乘以路由缩放因子\n",
        "        return weights.type_as(x), indices  # 返回计算后的权重和选择的专家索引\n",
        "\n",
        "\n",
        "class Expert(nn.Module):\n",
        "    \"\"\"\n",
        "    专家（Expert）层，用于 Mixture-of-Experts（MoE）模型。\n",
        "\n",
        "    该模块实现了一个独立的专家网络，每个专家由三层线性变换层组成。\n",
        "\n",
        "    属性:\n",
        "        w1 (nn.Module): 线性层，从输入到隐藏层的变换。\n",
        "        w2 (nn.Module): 线性层，从隐藏层到输出的变换。\n",
        "        w3 (nn.Module): 额外的线性层，用于特征变换。\n",
        "    \"\"\"\n",
        "    def __init__(self, dim: int, inter_dim: int):\n",
        "        \"\"\"\n",
        "        初始化 Expert 层。\n",
        "\n",
        "        参数:\n",
        "            dim (int): 输入和输出的维度。\n",
        "            inter_dim (int): 隐藏层的维度。\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.w1 = nn.Linear(dim, inter_dim)  # 输入到隐藏层的线性变换\n",
        "        self.w2 = nn.Linear(inter_dim, dim)  # 隐藏层到输出层的线性变换\n",
        "        self.w3 = nn.Linear(dim, inter_dim)  # 额外的线性变换层\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        Expert 层的前向计算。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量，形状为 (batch_size, dim)。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 经过 Expert 计算后的输出张量，形状为 (batch_size, dim)。\n",
        "        \"\"\"\n",
        "        return self.w2(F.silu(self.w1(x)) * self.w3(x))  # 先经过 w1 进行非线性变换，再与 w3 计算的结果相乘，最后通过 w2 输出\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "9fd81296b7f4974b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9fd81296b7f4974b",
        "outputId": "4dcc7c62-ee2e-48d4-a0e5-4a91c9d09e93"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "q shape: torch.Size([2, 128, 16, 192])\n",
            "q_nope shape: torch.Size([2, 128, 16, 128]), q_pe shape: torch.Size([2, 128, 16, 64])\n",
            "kv shape: torch.Size([2, 128, 576])\n",
            "kv shape: torch.Size([2, 128, 512]), k_pe shape: torch.Size([2, 128, 64])\n",
            "MoE x shape: torch.Size([2, 128, 2048])\n",
            "Gate x shape: torch.Size([256, 2048])\n",
            "Gate indices shape: torch.Size([256, 6])\n",
            "Gate weights : tensor([[0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        ...,\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625],\n",
            "        [0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625]], device='cuda:0',\n",
            "       dtype=torch.float32)\n",
            "torch.Size([2, 102400])\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.distributed as dist\n",
        "from typing import Optional\n",
        "\n",
        "class MoE(nn.Module):\n",
        "    \"\"\"\n",
        "    MoE（Mixture-of-Experts，专家混合）模块，用于选择性地激活多个专家网络，提高计算效率。\n",
        "\n",
        "    主要属性：\n",
        "        dim (int): 输入特征的维度。\n",
        "        n_routed_experts (int): 该模型中的总专家数量。\n",
        "        n_local_experts (int): 在分布式环境中，每个设备处理的专家数量。\n",
        "        n_activated_experts (int): 每个输入激活的专家数量。\n",
        "        gate (nn.Module): 用于计算输入到各专家的分配权重的门控机制。\n",
        "        experts (nn.ModuleList): 专家网络列表，每个专家都是一个神经网络模块。\n",
        "        shared_experts (nn.Module): 共享专家网络，对所有输入均生效。\n",
        "    \"\"\"\n",
        "    def __init__(self, args: ModelArgs):\n",
        "        \"\"\"\n",
        "        初始化 MoE 模块。\n",
        "\n",
        "        参数：\n",
        "            args (ModelArgs): 包含 MoE 参数的模型配置。\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.dim = args.dim\n",
        "\n",
        "        # 确保专家数量可以被世界大小整除（用于分布式训练）。\n",
        "        assert args.n_routed_experts % world_size == 0, f\"专家数量必须被 world_size 整除 (world_size={world_size})\"\n",
        "\n",
        "        self.n_routed_experts = args.n_routed_experts\n",
        "        self.n_local_experts = args.n_routed_experts // world_size\n",
        "        self.n_activated_experts = args.n_activated_experts\n",
        "\n",
        "        # 计算当前设备负责的专家索引范围。\n",
        "        self.experts_start_idx = rank * self.n_local_experts\n",
        "        self.experts_end_idx = self.experts_start_idx + self.n_local_experts\n",
        "\n",
        "        # 门控机制，用于决定输入数据分配给哪些专家。\n",
        "        self.gate = Gate(args)\n",
        "\n",
        "        # 仅在当前设备上初始化其负责的专家，其余设为 None。\n",
        "        self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None\n",
        "                                      for i in range(self.n_routed_experts)])\n",
        "\n",
        "        # 共享专家，对所有输入均适用。\n",
        "        self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        MoE 模块的前向传播。\n",
        "\n",
        "        参数：\n",
        "            x (torch.Tensor): 输入张量。\n",
        "\n",
        "        返回：\n",
        "            torch.Tensor: 经过专家计算后的输出张量。\n",
        "        \"\"\"\n",
        "        shape = x.size()\n",
        "        #打印x的shape\n",
        "        print(f\"MoE x shape: {x.shape}\")\n",
        "        x = x.view(-1, self.dim) #把x的bs和seq_len展平了\n",
        "\n",
        "        # 通过门控机制获取专家索引及其权重。\n",
        "        weights, indices = self.gate(x)\n",
        "        y = torch.zeros_like(x)\n",
        "\n",
        "        # 统计各专家被选中的次数。\n",
        "        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()\n",
        "\n",
        "        # 遍历当前设备管理的专家，并处理其对应的输入。\n",
        "        for i in range(self.experts_start_idx, self.experts_end_idx):\n",
        "            if counts[i] == 0:\n",
        "                continue  # 如果当前专家未被选中，则跳过。\n",
        "            expert = self.experts[i]\n",
        "            idx, top = torch.where(indices == i)\n",
        "            print(f'moe-idx{idx}')\n",
        "            y[idx] += expert(x[idx]) * weights[idx, top, None]\n",
        "\n",
        "        # 共享专家计算。\n",
        "        z = self.shared_experts(x)\n",
        "\n",
        "        # 若为多设备环境，则进行全局同步。\n",
        "        if world_size > 1:\n",
        "            dist.all_reduce(y)\n",
        "        #打印，y，z的shape\n",
        "        print(f\"MoE y shape: {y.shape}, MoE z shape: {z.shape}\")\n",
        "        return (y + z).view(shape)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.distributed as dist\n",
        "from typing import Optional\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\"\n",
        "    Transformer 块，结合了注意力机制和前馈网络。\n",
        "\n",
        "    属性:\n",
        "        attn (nn.Module): 多头注意力（MLA, Multi-Head Attention）。\n",
        "        ffn (nn.Module): 前馈神经网络（MLP 或 MoE）。\n",
        "        attn_norm (nn.Module): 注意力层的归一化层。\n",
        "        ffn_norm (nn.Module): 前馈网络的归一化层。\n",
        "    \"\"\"\n",
        "    def __init__(self, layer_id: int, args: ModelArgs):\n",
        "        \"\"\"\n",
        "        初始化 Transformer 块。\n",
        "\n",
        "        参数:\n",
        "            layer_id (int): 该层在 Transformer 中的索引。\n",
        "            args (ModelArgs): 包含 Transformer 相关参数的配置对象。\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.attn = MLA(args)  # 多头注意力机制\n",
        "        # 如果 layer_id 小于稠密层数量，则使用 MLP，否则使用 MoE 结构\n",
        "        self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)\n",
        "        self.attn_norm = RMSNorm(args.dim)  # 注意力归一化层\n",
        "        self.ffn_norm = RMSNorm(args.dim)   # 前馈归一化层\n",
        "\n",
        "    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        Transformer 块的前向传播。\n",
        "\n",
        "        参数:\n",
        "            x (torch.Tensor): 输入张量。\n",
        "            start_pos (int): 序列的起始位置。\n",
        "            freqs_cis (torch.Tensor): 预计算的旋转嵌入复指数值。\n",
        "            mask (Optional[torch.Tensor]): 掩码张量，用于排除特定位置的注意力。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 经过 Transformer 块计算后的输出张量。\n",
        "        \"\"\"\n",
        "        x = x + self.attn(self.attn_norm(x), start_pos, freqs_cis, mask)  # 归一化后进行注意力计算并残差连接\n",
        "        x = x + self.ffn(self.ffn_norm(x))  # 归一化后进入前馈网络并残差连接\n",
        "        return x\n",
        "\n",
        "\n",
        "class Transformer(nn.Module):\n",
        "    \"\"\"\n",
        "    Transformer 模型，包括嵌入层、多层 Transformer 块、最终归一化层和输出层。\n",
        "\n",
        "    属性:\n",
        "        max_seq_len (int): 最大序列长度。\n",
        "        embed (nn.Module): 词嵌入层。\n",
        "        layers (torch.nn.ModuleList): Transformer 块的列表。\n",
        "        norm (nn.Module): 所有 Transformer 层之后的归一化层。\n",
        "        head (nn.Module): 输出投影层，将隐藏状态映射到词汇表大小。\n",
        "        freqs_cis (torch.Tensor): 预计算的旋转嵌入复指数值。\n",
        "    \"\"\"\n",
        "    def __init__(self, args: ModelArgs):\n",
        "        \"\"\"\n",
        "        初始化 Transformer 模型。\n",
        "\n",
        "        参数:\n",
        "            args (ModelArgs): 包含 Transformer 相关参数的配置对象。\n",
        "        \"\"\"\n",
        "        global world_size, rank\n",
        "        world_size = dist.get_world_size() if dist.is_initialized() else 1  # 获取分布式训练的总进程数\n",
        "        rank = dist.get_rank() if dist.is_initialized() else 0  # 获取当前进程的 rank 值\n",
        "        Linear.dtype = torch.float8_e4m3fn if args.dtype == \"fp8\" else torch.bfloat16  # 设置默认数据类型\n",
        "        super().__init__()\n",
        "        self.max_seq_len = args.max_seq_len  # 最大序列长度\n",
        "        self.embed = ParallelEmbedding(args.vocab_size, args.dim)  # 词嵌入层，支持并行计算\n",
        "        self.layers = torch.nn.ModuleList()\n",
        "        for layer_id in range(args.n_layers):\n",
        "            self.layers.append(Block(layer_id, args))  # 添加多个 Transformer 块\n",
        "        self.norm = RMSNorm(args.dim)  # 归一化层\n",
        "        self.head = ColumnParallelLinear(args.dim, args.vocab_size, dtype=torch.get_default_dtype())  # 输出投影层\n",
        "        self.register_buffer(\"freqs_cis\", precompute_freqs_cis(args), persistent=False)  # 预计算旋转位置编码\n",
        "\n",
        "    @torch.inference_mode()\n",
        "    def forward(self, tokens: torch.Tensor, start_pos: int = 0):\n",
        "        \"\"\"\n",
        "        Transformer 的前向传播。\n",
        "\n",
        "        参数:\n",
        "            tokens (torch.Tensor): 形状为 (batch_size, seq_len) 的输入 token ID。\n",
        "            start_pos (int, 可选): 序列的起始位置，默认为 0。\n",
        "\n",
        "        返回:\n",
        "            torch.Tensor: 形状为 (batch_size, vocab_size) 的 logits。\n",
        "        \"\"\"\n",
        "        seqlen = tokens.size(1)  # 获取输入序列长度\n",
        "        h = self.embed(tokens)  # 通过词嵌入层获取 token 表示\n",
        "        freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]  # 获取对应位置的旋转位置编码\n",
        "        mask = None\n",
        "        if seqlen > 1:\n",
        "            mask = torch.full((seqlen, seqlen), float(\"-inf\"), device=tokens.device).triu_(1)  # 构造上三角掩码（防止未来信息泄露）\n",
        "        for layer in self.layers:\n",
        "            h = layer(h, start_pos, freqs_cis, mask)  # 依次通过每个 Transformer 块\n",
        "        h = self.norm(h)[:, -1]  # 归一化后取最后一个时间步的输出\n",
        "        logits = self.head(h)  # 通过输出投影层计算 logits\n",
        "        if world_size > 1:\n",
        "            all_logits = [torch.empty_like(logits) for _ in range(world_size)]\n",
        "            dist.all_gather(all_logits, logits)  # 在所有进程间收集 logits\n",
        "            logits = torch.cat(all_logits, dim=-1)  # 拼接所有进程的 logits\n",
        "        return logits\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    torch.set_default_dtype(torch.bfloat16)  # 设置默认数据类型\n",
        "    torch.set_default_device(\"cuda\")  # 设置默认计算设备为 GPU\n",
        "    torch.manual_seed(0)  # 设置随机种子，保证可复现性\n",
        "    args = ModelArgs()  # 初始化模型参数\n",
        "    x = torch.randint(0, args.vocab_size, (2, 128))  # 随机生成 token ID 作为输入\n",
        "    model = Transformer(args)  # 初始化 Transformer 模型\n",
        "    print(model(x).size())  # 运行模型并打印输出张量的形状，为啥只有一个词，因为只取了最后一个的输出，可以理解这个样例像咱们的bert基座一样，让大家去练习一个分类问题\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.3"
    },
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "L4"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}